{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "30f2ec3d-7408-415e-ba76-60202c602df2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sns.set(style=\"ticks\")\n",
    "np.set_printoptions(precision=1, suppress=True, threshold=5)\n",
    "pd.set_option('display.precision', 3)\n",
    "pd.set_option('display.max_rows', None)  # Show all rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c7669a15-1dcf-4129-b6ae-e7b2bcc3ad8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(patterns, metric):\n",
    "    dfs = []\n",
    "    for pattern in patterns:\n",
    "        files = glob.glob(pattern)\n",
    "        files.sort()\n",
    "        for fname in files:\n",
    "            setting = fname.split('/eval_')[1].replace('.json', '').split('_')\n",
    "\n",
    "            df = pd.read_json(fname, orient='index').transpose()\n",
    "            df['seed'] = int(setting[0])\n",
    "            df['runs'] = int(setting[1])\n",
    "            df['noise'] = float(setting[2])\n",
    "\n",
    "            df['method'] = setting[3]\n",
    "            horizon = setting[4].split('-')\n",
    "            df['ph'] = int(horizon[0])\n",
    "            df['ah'] = int(horizon[1])\n",
    "\n",
    "            # Set default values\n",
    "            df['nsample'] = 1\n",
    "            df['nmode'] = 1\n",
    "            df['decay'] = np.nan\n",
    "\n",
    "            if setting[3] == 'coherence':\n",
    "                df['nsample'] = int(setting[5])\n",
    "                df['decay'] = float(setting[6])\n",
    "            elif setting[3] in ['positive']:\n",
    "                df['nsample'] = int(setting[5])\n",
    "                df['nmode'] = int(setting[6])\n",
    "            elif setting[3] in ['contrast', 'positive', 'negative']:\n",
    "                df['nsample'] = int(setting[5])\n",
    "                df['nmode'] = int(setting[6])\n",
    "            elif setting[3] == 'bid':\n",
    "                df['nsample'] = int(setting[5])\n",
    "                df['nmode'] = int(setting[6])\n",
    "                df['decay'] = float(setting[7])\n",
    "            elif setting[3] in ['cma', 'cwarm']:\n",
    "                df['nsample'] = int(setting[5])\n",
    "                df['decay'] = float(setting[6])\n",
    "            elif setting[3] == 'ema':\n",
    "                df['decay'] = float(setting[5])\n",
    "\n",
    "            dfs.append(df)\n",
    "    \n",
    "    dff = pd.concat(dfs, ignore_index=True)\n",
    "    dff['decay'] = dff['decay'].astype(float)\n",
    "    dff['nsample'] = dff['nsample'].astype(int)\n",
    "    dff['nmode'] = dff['nmode'].astype(int)\n",
    "    dff['runs'] = dff['runs'].astype(int)\n",
    "    dff[metric] = dff[metric].astype(float)\n",
    "    \n",
    "    method_order = {'random': 0, 'warmstart': 0.8, 'ema': 1, 'bid': 2}\n",
    "    dff['sort_key'] = dff['method'].apply(lambda x: method_order.get(x.lower(), len(method_order)))\n",
    "    \n",
    "    dff = dff.sort_values(['noise', 'ah', 'sort_key', 'method', 'nsample', 'decay', 'nmode', 'seed', metric]).drop('sort_key', axis=1)\n",
    "    cols = ['noise', 'ah', 'method', 'nsample', 'decay', 'nmode', 'seed', metric]\n",
    "\n",
    "    dff = dff[cols]\n",
    "\n",
    "    return dff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "92947369-e83f-43a9-8616-f6c191fb199f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>noise</th>\n",
       "      <th>ah</th>\n",
       "      <th>method</th>\n",
       "      <th>nsample</th>\n",
       "      <th>decay</th>\n",
       "      <th>nmode</th>\n",
       "      <th>seed</th>\n",
       "      <th>test/mean_score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.0</td>\n",
       "      <td>1</td>\n",
       "      <td>random</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.846</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.0</td>\n",
       "      <td>8</td>\n",
       "      <td>random</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.884</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.0</td>\n",
       "      <td>8</td>\n",
       "      <td>warmstart</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.887</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>0.0</td>\n",
       "      <td>8</td>\n",
       "      <td>ema</td>\n",
       "      <td>1</td>\n",
       "      <td>0.5</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.866</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>0.0</td>\n",
       "      <td>8</td>\n",
       "      <td>bid</td>\n",
       "      <td>15</td>\n",
       "      <td>0.5</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>0.928</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>random</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.805</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>warmstart</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.852</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>ema</td>\n",
       "      <td>1</td>\n",
       "      <td>0.5</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.823</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1.0</td>\n",
       "      <td>1</td>\n",
       "      <td>bid</td>\n",
       "      <td>15</td>\n",
       "      <td>0.5</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>0.889</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1.0</td>\n",
       "      <td>8</td>\n",
       "      <td>random</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.582</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   noise  ah     method  nsample  decay  nmode  seed  test/mean_score\n",
       "0    0.0   1     random        1    NaN      1     0            0.846\n",
       "1    0.0   8     random        1    NaN      1     0            0.884\n",
       "4    0.0   8  warmstart        1    NaN      1     0            0.887\n",
       "6    0.0   8        ema        1    0.5      1     0            0.866\n",
       "8    0.0   8        bid       15    0.5      3     0            0.928\n",
       "2    1.0   1     random        1    NaN      1     0            0.805\n",
       "5    1.0   1  warmstart        1    NaN      1     0            0.852\n",
       "7    1.0   1        ema        1    0.5      1     0            0.823\n",
       "9    1.0   1        bid       15    0.5      3     0            0.889\n",
       "3    1.0   8     random        1    NaN      1     0            0.582"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Define tasks and corresponding metrics\n",
    "task = 'pusht'\n",
    "metric = 'test/mean_score'\n",
    "foldername = 'outputs'\n",
    "\n",
    "# Loop through each task and metric\n",
    "patterns = [\n",
    "    f'../{foldername}/{task}/*/*/*/eval_*_random_*.json',\n",
    "    f'../{foldername}/{task}/*/*/*/eval_*_warmstart_*.json',\n",
    "    f'../{foldername}/{task}/*/*/*/eval_*_ema_*.json',\n",
    "    f'../{foldername}/{task}/*/*/*/eval_*_contrast_*.json',\n",
    "    f'../{foldername}/{task}/*/*/*/eval_*_bid_*.json',\n",
    "]\n",
    "df = load_data(patterns, metric)\n",
    "\n",
    "display(df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6909544-0187-4756-a82f-a45843e9b458",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
